#pragma once

#include "TinySHA1.hpp"
#include "base64.hpp"

#include "server.hpp"
#include "websocket.hpp"

class websocket_server;

/**
 * @brief A ws_server_connection contains the state of a websocket connection
 */
class websocket_connection {
    friend class websocket_server;

    static constexpr auto WEBSOCKET_SERVER_MAGIC_KEY = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

public:
    websocket_connection(websocket_server& serv, win::socket sock, const win::sockaddr_in& address)
        : sock(sock)
        , address(address)
        , server_instance(serv)
        , _recv_buffer([this](websocket_frame frame) { return _handle_websocket_frame(std::move(frame)); }, true) {
    }

    websocket_connection(const websocket_connection&) = delete;
    websocket_connection& operator=(const websocket_connection&) = delete;
    websocket_connection(websocket_connection&&) = delete;
    websocket_connection& operator=(websocket_connection&&) = delete;

    ~websocket_connection();

    win::socket sock;
    win::sockaddr_in address;
    websocket_server& server_instance;
    websocket_state state{websocket_state::CONNECTING};

    [[nodiscard]] bool available() const {
        return state == websocket_state::OPEN;
    }

private:
    /**
     * @brief to handle the received data and websocket frames
     *
     * @param buf recv buffer
     * @return true data handling success
     * @return false the connection should close after this
     */
    [[nodiscard]] bool _recv(std::vector<char> buf);

    std::chrono::time_point<std::chrono::system_clock, std::chrono::milliseconds> _last_recv_time;

    websocket_receive_buffer _recv_buffer;
    std::string _handshake_buffer;

    [[nodiscard]] bool _send_pong(std::vector<char> buf) const;
    [[nodiscard]] bool _send_handshake_response(const std::string& key) const;

    [[nodiscard]] bool _handle_websocket_frame(websocket_frame frame);
};

class websocket_server {
    friend class websocket_connection;

    static constexpr bool ENABLE_HEARTBEAT = true;
    static constexpr auto HEARTBEAT_INTERVAL = std::chrono::seconds(10);

public:
    websocket_server(const websocket_server& other) = delete;
    websocket_server(websocket_server&& other) noexcept = delete;
    websocket_server& operator=(const websocket_server& other) = delete;
    websocket_server& operator=(websocket_server&& other) noexcept = delete;

    using recv_hook_t = std::function<void(win::socket, const win::sockaddr_in&, websocket_frame)>;
    using close_hook_t = std::function<void(win::socket, const win::sockaddr_in&)>;
    using accept_hook_t = std::function<void(websocket_connection&)>;

    explicit websocket_server(u_short port = 80)
        : _server(port) {
        _initialize_socket_server();
    }

    ~websocket_server() {
        if (!_connections.empty()) {
            stop();
        }
        if constexpr (ENABLE_HEARTBEAT) {
            _heartbeat_exit = true;
            if (std::this_thread::get_id() != _heartbeat_thread.get_id() && _heartbeat_thread.joinable()) {
                {
                    std::lock_guard<std::mutex> lock(_heartbeat_mutex);
                    _heartbeat_cv.notify_one();
                }
                _heartbeat_thread.join();
            }
        }
    }

    void start();

    void register_recv_hook(const recv_hook_t& recv_hook) {
        _recv_hook = recv_hook;
    }

    void register_accept_hook(const accept_hook_t& accept_hook) {
        _accept_hook = accept_hook;
    }

    void register_close_hook(const close_hook_t& close_hook) {
        _close_hook = close_hook;
    }

    bool send(win::socket sock, std::vector<char> buf, websocket_frame_code opcode = TEXT_FRAME, bool FIN = true,
        const std::function<void()>& callback = nullptr);

    bool close(win::socket sock, int close_cause = 1000);

    bool stop();

private:
    websocket_connection& _add_new_connection(win::socket sock, const win::sockaddr_in& address);
    [[nodiscard]] websocket_connection* _get_connection(win::socket sock);
    [[nodiscard]] bool _remove_connection(win::socket sock);

    void _initialize_socket_server();

    std::mutex _m;
    std::map<win::socket, std::unique_ptr<websocket_connection>> _connections;

    std::atomic<bool> _heartbeat_exit{false};
    std::mutex _heartbeat_mutex;
    std::condition_variable _heartbeat_cv;
    std::thread _heartbeat_thread;

    recv_hook_t _recv_hook;
    accept_hook_t _accept_hook;
    close_hook_t _close_hook;

    socket_server _server;
};
